// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <melon/executors/threaded_repeating_function_runner.h>
#include <pollux/common/caching/async_data_cache.h>
#include <pollux/common/caching/ssd_file.h>
#include <pollux/common/memory/memory_arbitrator.h>

namespace melon {
    class CPUThreadPoolExecutor;
}

namespace kumo::pollux {
    namespace memory {
        class MemoryAllocator;
    }

    namespace cache {
        class AsyncDataCache;
    }

    /// Manages a background daemon thread to report stats through 'StatsReporter'.
    class PeriodicStatsReporter {
    public:
        struct Options {
            Options() {
            }

            const pollux::memory::MemoryAllocator *allocator{nullptr};
            uint64_t allocatorStatsIntervalMs{2'000};

            const pollux::cache::AsyncDataCache *cache{nullptr};
            uint64_t cacheStatsIntervalMs{60'000};

            const memory::MemoryArbitrator *arbitrator{nullptr};
            uint64_t arbitratorStatsIntervalMs{60'000};

            const memory::MemoryPool *spillMemoryPool{nullptr};
            uint64_t spillStatsIntervalMs{60'000};

            std::string toString() const {
                return fmt::format(
                    "allocatorStatsIntervalMs:{}, cacheStatsIntervalMs:{}, "
                    "arbitratorStatsIntervalMs:{}, spillStatsIntervalMs:{}",
                    allocatorStatsIntervalMs,
                    cacheStatsIntervalMs,
                    arbitratorStatsIntervalMs,
                    spillStatsIntervalMs);
            }
        };

        PeriodicStatsReporter(const Options &options = Options());

        /// Invoked to start the report daemon in background.
        void start();

        /// Invoked to stop the report daemon in background.
        void stop();

    private:
        // Add a task to run periodically.
        template<typename TFunc>
        void addTask(const std::string &taskName, TFunc &&func, size_t intervalMs) {
            scheduler_.add(
                taskName,
                [taskName,
                    intervalMs,
                    func = std::forward<TFunc>(func)]() mutable noexcept {
                    try {
                        func();
                    } catch (const std::exception &e) {
                        KLOG(ERROR) << "Error running periodic task " << taskName << ": "
                                << e.what();
                    }
                    return std::chrono::milliseconds(intervalMs);
                });
        }

        void reportCacheStats();

        void reportAllocatorStats();

        void reportArbitratorStats();

        void reportSpillStats();

        const pollux::memory::MemoryAllocator *const allocator_{nullptr};
        const pollux::cache::AsyncDataCache *const cache_{nullptr};
        const pollux::memory::MemoryArbitrator *const arbitrator_{nullptr};
        const pollux::memory::MemoryPool *const spillMemoryPool_{nullptr};
        const Options options_;

        cache::CacheStats lastCacheStats_;

        melon::ThreadedRepeatingFunctionRunner scheduler_;
    };

    /// Initializes and starts the process-wide periodic stats reporter. Before
    /// 'stopPeriodicStatsReporter()' is called, this method can only be called once
    /// process-wide, and additional calls to this method will throw.
    void startPeriodicStatsReporter(const PeriodicStatsReporter::Options &options);

    /// Stops the process-wide periodic stats reporter.
    void stopPeriodicStatsReporter();
} // namespace kumo::pollux
